# 最终版：从原始HDF文件出发，集“CloudSat预处理->时空匹配->剖面图绘制”于一体的自动化流水线
# 批量绘制剖面图，将cloudsat数据与风云数据进行时空匹配，对于匹配点超过200的数据进行绘图，绘图时过滤雷达反射率低于-20的值，但不影响云底高度的计算
import os
import numpy as np
import pandas as pd
import netCDF4 as nc
from scipy.spatial import cKDTree
from datetime import datetime, timedelta
import matplotlib.pyplot as plt
from tqdm import tqdm
import logging
import glob
import concurrent.futures
from functools import partial

# --- 导入您自己的CloudSat读取模块 ---
from read_CloudSat import reader

# --- 导入剖面图相关库 ---
from mpl_toolkits.axes_grid1 import make_axes_locatable
import cartopy.crs as ccrs
from cartopy.mpl.ticker import LongitudeFormatter, LatitudeFormatter

# --- 全局常量定义 (请根据您的实际情况修改) ---

# 1. 输入数据目录
# CloudSat原始HDF文件存放目录
GEOPROF_DIR = '/mnt/datastore/liudddata/cloudsat_data/cloudsat_lida/cloudsat_GEOPROF20200506'
LIDA_DIR = '/mnt/datastore/liudddata/cloudsat_data/cloudsat_lida/cloudsat_lida20200506'
# FY-4A 反演结果文件存放目录
FY_ROOT_DIR = "/mnt/datastore/liudddata/result/20200506_droupout"
# FY-4A 静态地理坐标文件
COORD_FILE = "FY4A_coordinates.nc"

# 2. 输出路径
# 所有剖面图的保存目录
PLOT_SAVE_DIR = '/mnt/datastore/liudddata/cloudsat_data/cloudsat_cbh_csv/final_profile_charts_2020_0506'

# 3. 处理参数
# 匹配点数超过多少才进行绘图
MIN_MATCHES_FOR_PLOT = 200
# MIN_MATCHES_FOR_PLOT = 100
# 使用的并行进程数
NUM_WORKERS = 8

# --- 日志和绘图设置 ---
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
plt.rcParams['figure.dpi'] = 200


# --- 函数定义部分 ---

# ==============================================================================
# 1. 核心处理函数 - 处理单个CloudSat文件
# ==============================================================================

def process_cloudsat_granule(file_pair, fy_kdtree_data, fy_root_dir, save_dir):
    """
    处理单个CloudSat文件对（GEOPROF和LIDA）的完整流程：
    读取 -> 计算CBH/CTH -> 匹配FY数据 -> 筛选 -> 绘制剖面图
    """
    geoprof_path, lida_path = file_pair
    try:
        # --- 步骤1: 一次性读取CloudSat所有需要的数据 ---
        f = reader(geoprof_path)
        d = reader(lida_path)

        # a. 用于绘图的原始数据
        lon_cs_raw, lat_cs_raw, elv_cs_raw = f.read_geo()
        height_cs_raw = f.read_sds('Height') / 1000.0  # 转为km
        reflectivity_cs_raw = f.read_sds('Radar_Reflectivity')

        # b. 用于匹配的计算数据
        cloud_mask = f.read_sds('CPR_Cloud_mask')
        cloud_faction = d.read_sds('CloudFraction')
        time_cs_raw = f.read_time(datetime=True)
        f.close()
        d.close()

        # --- 步骤2: 计算CloudSat的CTH/CBH (使用您提供的逻辑) ---
        cloudsat_cth, cloudsat_cbh = [], []
        for i in range(len(time_cs_raw)):
            valid_mask = (cloud_mask[i, :] >= 20) & (reflectivity_cs_raw[i, :] >= -30)
            # valid_mask = (cloud_mask[i, :] >= 20) & (reflectivity_cs_raw[i, :] >= -30)
            if np.any(valid_mask):
                valid_indices = np.where(valid_mask)[0]
                cloudsat_cth.append(height_cs_raw[i, valid_indices[0]])
                cloudsat_cbh.append(height_cs_raw[i, valid_indices[-1]])
            else:
                cloudsat_cth.append(np.nan)
                cloudsat_cbh.append(np.nan)

        # 创建副本用于绘图
        reflectivity_for_plot = reflectivity_cs_raw.copy()
        reflectivity_for_plot[reflectivity_for_plot <= -20] = np.nan

        # 将计算结果整理成DataFrame，用于后续匹配
        cloudsat_df = pd.DataFrame({
            'time': time_cs_raw,
            'longitude': lon_cs_raw,
            'latitude': lat_cs_raw,
            'cloudsat_cth': np.array(cloudsat_cth) * 1000,  # 转为m
            'cloudsat_cbh': np.array(cloudsat_cbh) * 1000,  # 转为m
        }).dropna(subset=['longitude', 'latitude'])

        if cloudsat_df.empty:
            return

        # --- 步骤3: 按小时对CloudSat轨道数据分段，并与FY数据匹配 ---
        cloudsat_df['hour_group'] = cloudsat_df['time'].dt.floor('H')
        for hour, event_group in cloudsat_df.groupby('hour_group'):

            # a. 寻找对应的FY文件
            fy_filename = os.path.join(fy_root_dir, f"{hour.strftime('%Y%m%d%H')}_predicted_2d_mc.nc")
            if not os.path.exists(fy_filename):
                continue

            # b. 读取FY数据
            fy_cth_data, fy_cbh_data, fy_clt_data = read_fy_data(fy_filename)

            # c. 执行时空匹配
            fy_tree, valid_mask, lat_fy, lon_fy = fy_kdtree_data
            matches = match_observations(
                fy_tree, valid_mask, lat_fy, lon_fy, fy_cth_data, fy_cbh_data, fy_clt_data,
                event_group['longitude'].values, event_group['latitude'].values,
                event_group['time'].values, event_group['cloudsat_cth'].values,
                event_group['cloudsat_cbh'].values
            )

            if not matches:
                continue

            # d. 过滤匹配结果
            matched_df = pd.DataFrame(matches)
            valid_match_mask = matched_df['cloudsat_cbh'].notna() & matched_df['fy_cbh'].notna()
            filtered_matches = matched_df[valid_match_mask].copy()

            # --- 步骤4: 如果有效匹配点数足够，立即绘图 ---
            if len(filtered_matches) >= MIN_MATCHES_FOR_PLOT:
                print(
                    f"找到一个重点事件: {hour.strftime('%Y-%m-%d %H:00')}, 有效匹配点数: {len(filtered_matches)}. 开始绘图...")

                # 筛选原始雷达数据用于绘图
                time_series_cs = pd.to_datetime(time_cs_raw)
                event_mask = (time_series_cs >= hour) & (time_series_cs < hour + timedelta(hours=1))

                # plot_profile_for_event(
                #     filtered_matches,
                #     reflectivity_cs_raw[event_mask, :].T,
                #     height_cs_raw[event_mask, :].T,
                #     lat_cs_raw[event_mask],
                #     lon_cs_raw[event_mask],
                #     elv_cs_raw[event_mask] / 1000.0,  # 转为km
                #     save_dir
                # )
                # 使用过滤后的反射率数据绘制剖面图
                plot_profile_for_event(
                    filtered_matches,
                    reflectivity_for_plot[event_mask, :].T,  # 使用过滤后的数据
                    height_cs_raw[event_mask, :].T,
                    lat_cs_raw[event_mask],
                    lon_cs_raw[event_mask],
                    elv_cs_raw[event_mask] / 1000.0,  # 转为km
                    save_dir
                )

    except Exception as e:
        # 使得单个文件的错误不影响整个程序的运行
        logging.error(f"处理文件 {os.path.basename(geoprof_path)} 时出错: {e}")


# ==============================================================================
# 2. 辅助函数 (匹配与绘图)
# ==============================================================================
# (此处省略了 read_fy_data, read_coordinate_data, spherical_to_cartesian,
# build_fy_kdtree, match_observations, set_map, draw_track, draw_elevation 函数，
# 因为它们与您之前的版本相同)

def read_fy_data(filename):
    with nc.Dataset(filename, 'r') as f:
        cth = f.variables['cth'][:]
        cbh = f.variables['predicted_mean'][:]
        clt = f.variables['clt'][:]
    return cth, cbh, clt


def read_coordinate_data():
    with nc.Dataset(COORD_FILE, 'r') as f:
        lat = f.variables['lat'][:, :].T
        lon = f.variables['lon'][:, :].T
    return lat, lon


def spherical_to_cartesian(lon, lat):
    lon_rad = np.deg2rad(lon)
    lat_rad = np.deg2rad(90.0 - lat)
    x = np.sin(lat_rad) * np.cos(lon_rad)
    y = np.sin(lat_rad) * np.sin(lon_rad)
    z = np.cos(lat_rad)
    return np.column_stack([x, y, z])


def build_fy_kdtree():
    lat_fy, lon_fy = read_coordinate_data()
    valid_mask = ~np.isnan(lat_fy)
    coords_cart = spherical_to_cartesian(lon_fy[valid_mask], lat_fy[valid_mask])
    return cKDTree(coords_cart), valid_mask, lat_fy, lon_fy


def match_observations(tree, valid_mask, lat_fy, lon_fy, cth_data, cbh_data, clt_data,
                       cloudsat_lons, cloudsat_lats, cloudsat_times,
                       cloudsat_cth, cloudsat_cbh, radius_km=2.0):
    matches = []
    radius_rad = radius_km / 6371.0
    cs_points = spherical_to_cartesian(cloudsat_lons, cloudsat_lats)
    indices_list = tree.query_ball_point(cs_points, r=radius_rad)
    valid_indices_rows, valid_indices_cols = np.where(valid_mask)

    for i, indices in enumerate(indices_list):
        if not indices: continue
        fy_linear_idx = indices[0]
        fy_row = valid_indices_rows[fy_linear_idx]
        fy_col = valid_indices_cols[fy_linear_idx]

        matches.append({
            'cloudsat_time': cloudsat_times[i], 'cloudsat_lat': cloudsat_lats[i], 'cloudsat_lon': cloudsat_lons[i],
            'cloudsat_cth': cloudsat_cth[i], 'cloudsat_cbh': cloudsat_cbh[i],
            'fy_lat': lat_fy[fy_row, fy_col], 'fy_lon': lon_fy[fy_row, fy_col],
            'fy_cth': cth_data[fy_row, fy_col], 'fy_cbh': cbh_data[fy_row, fy_col],
            'fy_clt': clt_data[fy_row, fy_col]
        })
    return matches


def set_map(ax):
    ax.coastlines(lw=0.5)
    ax.set_global()
    ax.set_xticks(np.arange(-180, 181, 60), crs=ccrs.PlateCarree())
    ax.set_yticks(np.arange(-90, 91, 30), crs=ccrs.PlateCarree())
    ax.xaxis.set_major_formatter(LongitudeFormatter())
    ax.yaxis.set_major_formatter(LatitudeFormatter())


def draw_track(ax, lon1D, lat1D, title):
    ax.plot(lon1D, lat1D, lw=2, color='b', transform=ccrs.Geodetic())
    ax.plot(lon1D[0], lat1D[0], 'ro', ms=3, transform=ccrs.PlateCarree())
    ax.text(lon1D[0] + 5, lat1D[0], 'start', color='r', transform=ccrs.PlateCarree())
    ax.set_title(title, fontsize='medium')


def draw_cross_section(ax, x_coords_1d, hgt_1d, data_2d, fy_cth, fy_cbh, fy_x_coords):
    """画出data的剖面，并添加风云数据的云底和云顶高度折线图."""
    # 使用清洗后的一维坐标和二维数据进行绘图
    im = ax.pcolormesh(x_coords_1d, hgt_1d, data_2d, cmap='jet', shading='auto', vmin=-25, vmax=25)

    ax.set_ylim(0, 20)
    ax.set_xlim(x_coords_1d.min(), x_coords_1d.max())
    ax.set_xlabel('Latitude')
    ax.set_ylabel('Height [km]')

    # 绘制FY数据覆盖层
    ax.plot(fy_x_coords, fy_cth / 1000, color='red', label='FY CTH', linewidth=1)
    ax.plot(fy_x_coords, fy_cbh / 1000, color='yellow', label='FY CBH', linewidth=1, linestyle='--')
    ax.legend(loc='upper right')

    # 设置颜色条
    divider = make_axes_locatable(ax)
    cax = divider.append_axes('bottom', size='5%', pad=0.4)
    cbar = plt.colorbar(im, cax=cax, extend='both', orientation='horizontal')
    cbar.set_label('Radar Reflectivity [dBZe]')


def draw_elevation(ax, x_coords, elv):
    ax.fill_between(x_coords, elv, color='gray')


def plot_profile_for_event(matched_event_df, reflectivity_data, height_data, lat_track, lon_track, elev_track,
                           save_dir):
    """为单个匹配事件绘制剖面图 - 已增加对纬度和高度数据的双重清洗."""
    if matched_event_df.empty:
        return

    start_time = matched_event_df['cloudsat_time'].min()
    event_hour = start_time.replace(minute=0, second=0, microsecond=0)

    # --- 步骤 1: 清洗纬度坐标 (X轴) ---
    valid_coords_mask = np.isfinite(lat_track)
    if not np.any(valid_coords_mask):
        print(f"警告：事件 {start_time.strftime('%Y-%m-%d %H:%M')} 的所有CloudSat坐标点均无效，跳过绘图。")
        return

    # 应用掩码过滤所有与剖面维度相关的1D数组
    lat_t_clean = lat_track[valid_coords_mask]
    lon_t_clean = lon_track[valid_coords_mask]
    elv_t_clean = elev_track[valid_coords_mask]

    # 应用掩码过滤所有与剖面维度相关的2D数组
    # 注意: 数组的形状是 (层数, 剖面数)，所以掩码应用在第2个轴上
    rad_refl_t_clean = reflectivity_data[:, valid_coords_mask]
    hgt_2d_clean = height_data[:, valid_coords_mask]

    # --- 步骤 2: 清洗高度坐标 (Y轴) ---
    # 计算每个高度层的平均值，忽略NaN，得到一个干净的1D高度数组
    # 这是绘制pcolormesh更稳健的方式
    hgt_1d_final = np.nanmean(hgt_2d_clean, axis=1)

    # 再次检查生成的一维高度数组是否还有无效值
    if np.any(~np.isfinite(hgt_1d_final)):
        print(f"警告：事件 {start_time.strftime('%Y-%m-%d %H:%M')} 的高度数据在平均后仍包含无效值，跳过绘图。")
        return

    # 按纬度排序匹配数据，以确保FY折线图正确绘制
    plot_data = matched_event_df.sort_values(by='cloudsat_lat')

    # --- 步骤 3: 开始绘图 ---
    fig = plt.figure(dpi=200, figsize=(10, 10))
    ax1 = fig.add_axes([0.1, 0.55, 0.8, 0.4], projection=ccrs.PlateCarree())
    ax2 = fig.add_axes([0.1, 0.1, 0.8, 0.4])

    map_title = f"CloudSat Track for {event_hour.strftime('%Y-%m-%d %H:00 UTC')}"
    set_map(ax1)
    draw_track(ax1, lon_t_clean, lat_t_clean, map_title)

    # 将清洗后的所有数据传递给绘图函数
    draw_cross_section(ax2, lat_t_clean, hgt_1d_final, rad_refl_t_clean,
                       plot_data['fy_cth'].values, plot_data['fy_cbh'].values,
                       plot_data['fy_lat'].values)
    draw_elevation(ax2, lat_t_clean, elv_t_clean)
    ax2.set_title('Radar Reflectivity & FY-4A Cloud Boundaries')

    save_path = os.path.join(save_dir, f'profile_{event_hour.strftime("%Y%m%d_%H%M")}.png')
    fig.savefig(save_path, bbox_inches='tight')
    plt.close(fig)



# --- 主执行部分 (Main Execution Block) ---
if __name__ == "__main__":
    # 确保输出目录存在
    if not os.path.exists(PLOT_SAVE_DIR):
        os.makedirs(PLOT_SAVE_DIR)
        print(f"已创建目录: {PLOT_SAVE_DIR}")

    # 1. 一次性构建FY-4A的KD树
    print("--- 步骤 1: 正在构建风云4号坐标系的KD树 (只需一次)... ---")
    fy_kdtree_data = build_fy_kdtree()
    print("KD树构建完成。")

    # 2. 准备CloudSat文件列表 (使用更健壮的配对方法)
    print("--- 步骤 2: 正在准备CloudSat文件列表... ---")

    # --- 新增的健壮文件配对逻辑 ---
    geoprof_files_map = {}
    # 首先，为所有GEOPROF文件创建一个基于轨道号的查找字典
    for f in os.listdir(GEOPROF_DIR):
        if f.endswith('.hdf'):
            # 从文件名中提取轨道号作为唯一键，例如 '74925'
            try:
                granule_key = f.split('_')[1]
                geoprof_files_map[granule_key] = os.path.join(GEOPROF_DIR, f)
            except IndexError:
                print(f"警告：无法从文件名 {f} 中解析轨道号，已跳过。")
                continue

    file_pairs = []
    # 然后，遍历LIDA文件，去字典中寻找配对的GEOPROF文件
    for f in os.listdir(LIDA_DIR):
        if f.endswith('.hdf'):
            try:
                granule_key = f.split('_')[1]
                if granule_key in geoprof_files_map:
                    # 如果找到了配对，则将(GEOPROF路径, LIDA路径)加入列表
                    file_pairs.append((geoprof_files_map[granule_key], os.path.join(LIDA_DIR, f)))
                else:
                    print(f"警告：为LIDA文件 {f} 未找到匹配的GEOPROF文件，已跳过。")
            except IndexError:
                print(f"警告：无法从文件名 {f} 中解析轨道号，已跳过。")
                continue

    if not file_pairs:
        raise ValueError("未能找到任何成功配对的GEOPROF和LIDA文件，请检查目录和文件名格式！")

    print(f"共找到 {len(file_pairs)} 对可处理的CloudSat文件。")

    # 3. 使用多进程并行处理所有CloudSat文件
    print(f"--- 步骤 3: 开始使用 {NUM_WORKERS} 个进程并行处理... ---")

    task_function = partial(process_cloudsat_granule,
                            fy_kdtree_data=fy_kdtree_data,
                            fy_root_dir=FY_ROOT_DIR,
                            save_dir=PLOT_SAVE_DIR)

    with concurrent.futures.ProcessPoolExecutor(max_workers=NUM_WORKERS) as executor:
        list(tqdm(executor.map(task_function, file_pairs), total=len(file_pairs), desc="处理CloudSat文件"))

    print("\n所有任务已成功完成！ (All tasks completed successfully!)")